#!/usr/bin/env python3
"""
ELFUtils MCP Server

面向 elfutils 软件包的 MCP 服务器，提供以下工具：

1. elfutils_debuginfo_management - 调试信息文件管理
   用途：创建调试文件、剥离符号、添加调试链接
   参数：source_file, debug_file, stripped_file
   返回：统一JSON结构

2. elfutils_stack_analysis - 堆栈分析
   用途：分析核心转储文件，获取堆栈信息和源代码位置
   参数：executable, core_file, address
   返回：统一JSON结构

3. elfutils_service_control - 服务控制
   用途：控制 debuginfod 服务的启动、停止和重启
   参数：action (start/stop/restart)
   返回：统一JSON结构
"""

import subprocess
import shlex
from typing import List
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("ELFUtils MCP Server")

@mcp.tool()
def elfutils_debuginfo_management(source_file: str, debug_file: str, stripped_file: str) -> dict:
    """
    调试信息文件管理工具
    
    参数:
        source_file: 源可执行文件路径
        debug_file: 调试信息文件路径
        stripped_file: 剥离符号后的文件路径
    
    返回:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    commands = [
        f"objcopy --only-keep-debug {shlex.quote(source_file)} {shlex.quote(debug_file)}",
        f"strip {shlex.quote(source_file)}",
        f"objcopy --add-gnu-debuglink={shlex.quote(debug_file)} {shlex.quote(source_file)}",
        f"cp {shlex.quote(debug_file)} {shlex.quote(stripped_file)}",
        f"eu-unstrip {shlex.quote(source_file)} {shlex.quote(stripped_file)}"
    ]
    
    results = []
    for cmd in commands:
        try:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
            results.append({
                "success": result.returncode == 0,
                "command": cmd,
                "exit_code": result.returncode,
                "stdout": result.stdout,
                "stderr": result.stderr
            })
        except subprocess.TimeoutExpired:
            results.append({
                "success": False,
                "command": cmd,
                "exit_code": -1,
                "stdout": "",
                "stderr": "Command timed out"
            })
        except Exception as e:
            results.append({
                "success": False,
                "command": cmd,
                "exit_code": -1,
                "stdout": "",
                "stderr": str(e)
            })
    
    # 返回最后一个命令的结果作为主要结果
    return results[-1] if results else {
        "success": False,
        "command": "debuginfo_management",
        "exit_code": -1,
        "stdout": "",
        "stderr": "No commands executed"
    }

@mcp.tool()
def elfutils_stack_analysis(executable: str, core_file: str, address: str = None) -> dict:
    """
    堆栈分析工具
    
    参数:
        executable: 可执行文件路径（带调试信息）
        core_file: 核心转储文件路径
        address: 特定地址（可选），如未提供则自动从堆栈中获取
    
    返回:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    try:
        # 如果未提供地址，从堆栈中获取第一个地址
        if address is None:
            stack_cmd = f"eu-stack -e {shlex.quote(executable)} --core {shlex.quote(core_file)}"
            stack_result = subprocess.run(stack_cmd, shell=True, capture_output=True, text=True, timeout=30)
            
            if stack_result.returncode == 0 and stack_result.stdout:
                # 提取第一个地址（#0 后面的地址）
                for line in stack_result.stdout.split('\n'):
                    if line.startswith('#0'):
                        parts = line.split()
                        if len(parts) >= 2:
                            address = parts[1]
                            break
        
        if address:
            cmd = f"eu-addr2line -e {shlex.quote(executable)} --core {shlex.quote(core_file)} {shlex.quote(address)}"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
            return {
                "success": result.returncode == 0,
                "command": cmd,
                "exit_code": result.returncode,
                "stdout": result.stdout,
                "stderr": result.stderr
            }
        else:
            return {
                "success": False,
                "command": "address_extraction",
                "exit_code": -1,
                "stdout": "",
                "stderr": "Could not extract address from stack trace"
            }
            
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": "stack_analysis",
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": "stack_analysis",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def elfutils_service_control(action: str) -> dict:
    """
    调试信息服务控制工具
    
    参数:
        action: 服务操作，可选值：start, stop, restart
    
    返回:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    valid_actions = ['start', 'stop', 'restart']
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} debuginfod.service",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Must be one of {valid_actions}"
        }
    
    try:
        cmd = f"systemctl {shlex.quote(action)} debuginfod.service"
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": cmd,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()